#!/usr/bin/python
# -*-coding:utf-8-*-

import os
from inspect import isfunction
import re
from copy import deepcopy

import dill
import graphviz
import six.moves.cPickle as pickle

from base.utils import np_keys_to_pd
from base.utils import np_neu_keys_to_pd

from gplearn_evolve.base.GPFunctionsV1 import *
from gplearn_evolve.base.GPFunctionsV2 import *
from gplearn_evolve.base.GPFunctionsV3 import *
from gplearn_evolve.base.GPFunctionsCWV1 import *
from gplearn_evolve.base.GPFunctionsCWV2 import *

def save_gpmodel(program, path, base=None):
    if base is None:
        program = program._program.program
    else:
        program = [program._program.program, base]

    with open(path, 'wb') as f:
        pickle.dump(program, f)

def save_gpmodel_v2(program, path, base=None):
    if base is None:
        save_program = program
    else:
        save_program = [program, base]

    with open(path, 'wb') as f:
        pickle.dump(save_program, f)

def load_gpmodel(path):
    with open(path, 'rb') as f:
        program = pickle.load(f)

        return program

## TODO - predict
def _func_data_filter(X, apply_stack):
    terminals_list = []

    # function_name = apply_stack[-1][0].name
    for t in apply_stack[-1][1:]:
        if isinstance(t, int):
            terminals_list.append(X[:, t])
        elif isinstance(t, (float, np.float32, np.float64)):
            if t < 1.0:
                ###### ZBC: 这种情况一般是t是应该参数，例如ts_函数里面的天数d，其它基本没什么改动
                terminals_list.append(t)
            else:
                terminals_list.append(np.repeat(t, X.shape[0]))
        else:
            ###### ZBC: 这里是参入函数
            terminals_list.append(t)

    return terminals_list

def _func_data_filter_v2(X, apply_stack):
    terminals_list = []

    # function_name = apply_stack[-1][0].name
    for t in apply_stack[-1][1:]:
        if isinstance(t, str):
            terminals_list.append(X[t].values)
        elif isinstance(t, (float, np.float32, np.float64)):
            if t < 1.0:
                ###### ZBC: 这种情况一般是t是应该参数，例如ts_函数里面的天数d，其它基本没什么改动
                terminals_list.append(t)
            else:
                terminals_list.append(np.repeat(t, X.shape[0]))
        else:
            ###### ZBC: 这里是参入函数
            terminals_list.append(t)

    return terminals_list

def model_predict(X, program, keys=None, neu_keys=None, by_name=False, key_by_path=True):
    if len(program) > 1:
        base = program[1]
        gp_program = program[0]
    else:
        gp_program = program
        base = 5

    apply_stack = []
    for node in gp_program:
        if not isinstance(node, (str, int, float)):
            apply_stack.append([node])
        else:
            # Lazily evaluate later
            apply_stack[-1].append(node)

        while len(apply_stack[-1]) == apply_stack[-1][0].arity + 1:
            # Apply functions that have sufficient arguments
            function = apply_stack[-1][0]

            func_name = function.name

            if func_name.startswith('ts_') and (not func_name.startswith('ts_cw_')):
                used_func = eval('def_keyfunc_' + func_name)(keys, base=base, key_by_path=key_by_path)
            elif func_name.startswith('cs_') or func_name.startswith('ts_cw_'):
                if 'neu' in func_name:
                    used_func = eval('def_keyfunc_' + func_name)(keys, neu_keys, key_by_path=key_by_path)
                else:
                    used_func = eval('def_keyfunc_' + func_name)(keys, key_by_path=key_by_path)
            else:
                used_func = function

            ###### ZBC: 主要改动，见下面的函数
            if by_name:
                terminals = _func_data_filter_v2(X, apply_stack)
            else:
                terminals = _func_data_filter(X, apply_stack)
            # terminals = _func_data_filter(X, apply_stack)

            # TODO - run function
            intermediate_result = used_func(*terminals)

            if len(apply_stack) != 1:
                apply_stack.pop()
                apply_stack[-1].append(intermediate_result)
            else:
                return intermediate_result
                # print(intermediate_result)
                # break




